from fastapi import WebSocket, APIRouter
from typing import List

router = APIRouter()

# 存储客户端 ID 和 WebSocket 之间的映射关系
clients = {}

# 存储客户端 ID 和用户 ID 之间的映射关系
client_user_mapping = {}


# WebSocket 连接处理函数
async def handle_websocket(websocket: WebSocket, client_id: int, user_id: str):
    await websocket.accept()
    clients[client_id] = websocket
    client_user_mapping[client_id] = user_id
    try:
        while True:
            data = await websocket.receive_text()
            broadcast_message(f"User {user_id}: {data}", exclude=[client_id])
    finally:
        del clients[client_id]
        del client_user_mapping[client_id]


# 广播消息给所有客户端
def broadcast_message(message: str, exclude: List[int] = []):
    for client_id, client in clients.items():
        if client_id not in exclude:
            client.send_text(message)


# WebSocket 端点，参数中包含用户 ID
# @router.websocket("/ws/{user_id}")
# async def websocket_endpoint(websocket: WebSocket, user_id: str):
#     # 生成一个客户端 ID
#     client_id = id(websocket)
#     await handle_websocket(websocket, client_id, user_id)


@router.get("/")
async def read_root():
    return {"Hello": "World"}
